{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# HybridNonlinearFactor"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<a href=\"https://colab.research.google.com/github/borglab/gtsam/blob/develop/gtsam/hybrid/doc/HybridNonlinearFactor.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "tags": [
          "remove-cell"
        ]
      },
      "outputs": [],
      "source": [
        "try:\n",
        "    import google.colab\n",
        "    %pip install --quiet gtsam-develop\n",
        "except ImportError:\n",
        "    pass  # Not running on Colab, do nothing"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "A `HybridNonlinearFactor` represents a factor in a **nonlinear** hybrid optimization problem. Similar to `HybridGaussianFactor`, it represents a factor whose behavior depends on the assignment of discrete variables, but here the underlying components are `gtsam.NonlinearFactor`s (specifically, `NoiseModelFactor`s). It inherits from `HybridFactor`.\n",
        "\n",
        "Mathematically, it represents a factor $\\phi(X, M)$ whose value for a specific discrete assignment $M=m$ is proportional to the likelihood of the continuous variables $X$ under that mode, potentially weighted:\n",
        "$$\n",
        "\\phi(X, M=m) \\propto \\exp\\left(-E_m - \\frac{1}{2} ||h_m(X) - z_m||^2_{\\Sigma_m}\\right)\n",
        "$$\n",
        "This corresponds to an error (negative log-likelihood) for mode $m$ of:\n",
        "$$\n",
        "\\text{error}(X, M=m) = E_m + \\frac{1}{2} ||h_m(X) - z_m||^2_{\\Sigma_m}\n",
        "$$\n",
        "where for each discrete assignment $m$, we have:\n",
        "*   A nonlinear measurement function $h_m(X)$.\n",
        "*   A measurement $z_m$.\n",
        "*   A noise model $\\Sigma_m$ defining the squared Mahalanobis distance $|| \\cdot ||^2_{\\Sigma_m}$.\n",
        "*   An optional scalar term $E_m$. This term is added to the negative log-likelihood derived from the `NoiseModelFactor`. It can be used, for example, to incorporate prior probabilities $P(m)$ for different modes by setting $E_m = -\\log P(m)$ (plus any constant, as only differences in $E_m$ matter for optimization). A smaller $E_m$ makes mode $m$ relatively more likely, all else being equal.\n",
        "\n",
        "Internally, it uses a `DecisionTree<Key, NonlinearFactorValuePair>` to store the different nonlinear factor components (`NoiseModelFactor::shared_ptr`) and their associated scalar terms (`double`), indexed by the `DiscreteKey`s.\n",
        "\n",
        "```mermaid\n",
        "graph TD\n",
        "    HybridNonlinearFactor --> HybridFactor\n",
        "    HybridNonlinearFactor -- Uses --> DecisionTree[\"DecisionTree<Key, NonlinearFactorValuePair>\"]\n",
        "    DecisionTree -- Stores Leaf Values --> NonlinearFactorValuePair\n",
        "    NonlinearFactorValuePair -- Contains --> ComponentFactor[\"NoiseModelFactor::shared_ptr <br> (embodies h<sub>m</sub>, z<sub>m</sub>, Σ<sub>m</sub>)\"]\n",
        "    NonlinearFactorValuePair -- Contains --> ScalarTerm[\"double <br> (the E<sub>m</sub> scalar term)\"]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [],
      "source": [
        "import gtsam\n",
        "import numpy as np\n",
        "\n",
        "from gtsam import (\n",
        "    HybridNonlinearFactor,\n",
        "    BetweenFactorPose2,\n",
        "    Values,\n",
        "    DiscreteValues,\n",
        "    HybridValues,\n",
        "    Pose2)\n",
        "from gtsam.symbol_shorthand import X, D"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Initialization"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Initialize with discrete keys and corresponding `NoiseModelFactor` components, optionally with scalar energies."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "HybridNonlinearFactor (factors only):\n",
            "HybridNonlinearFactor\n",
            " Hybrid [x0 x1; d0]\n",
            "HybridNonlinearFactor\n",
            " Choice(d0) \n",
            " 0 Leaf (0) BetweenFactor(x0,x1)\n",
            "  measured:  (1, 0, 0)\n",
            "  noise model: diagonal sigmas [0.5; 0.5; 0.174532925];\n",
            "\n",
            " 1 Leaf (0) BetweenFactor(x0,x1)\n",
            "  measured:  (1, 0, 0)\n",
            "  noise model: diagonal sigmas [0.1; 0.1; 0.0174532925];\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# --- Example: Mode-Dependent Odometry ---\n",
        "dk0 = (D(0), 2) # Binary mode: Slippery (0) or Grippy (1)\n",
        "\n",
        "# Nonlinear factors (BetweenFactorPose2) for each mode\n",
        "# Mode 0 (Slippery): Larger noise\n",
        "noise0 = gtsam.noiseModel.Diagonal.Sigmas([0.5, 0.5, np.radians(10)])\n",
        "odom0 = BetweenFactorPose2(X(0), X(1), Pose2(1.0, 0, 0), noise0)\n",
        "\n",
        "# Mode 1 (Grippy): Smaller noise\n",
        "noise1 = gtsam.noiseModel.Diagonal.Sigmas([0.1, 0.1, np.radians(1)])\n",
        "odom1 = BetweenFactorPose2(X(0), X(1), Pose2(1.0, 0, 0), noise1)\n",
        "\n",
        "# Option A: Just factors (scalar energy E_m = 0)\n",
        "hnlf_a = HybridNonlinearFactor(dk0, [odom0, odom1])\n",
        "print(\"HybridNonlinearFactor (factors only):\")\n",
        "hnlf_a.print()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "id": "624e20b8",
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "HybridNonlinearFactor (factors + scalars):\n",
            "HybridNonlinearFactor\n",
            " Hybrid [x0 x1; d0]\n",
            "HybridNonlinearFactor\n",
            " Choice(d0) \n",
            " 0 Leaf (1.60943791) BetweenFactor(x0,x1)\n",
            "  measured:  (1, 0, 0)\n",
            "  noise model: diagonal sigmas [0.5; 0.5; 0.174532925];\n",
            "\n",
            " 1 Leaf (0.223143551) BetweenFactor(x0,x1)\n",
            "  measured:  (1, 0, 0)\n",
            "  noise model: diagonal sigmas [0.1; 0.1; 0.0174532925];\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# Option B: Factors + scalar energies (e.g., prior on mode)\n",
        "# P(Slippery)=0.2 -> E0 = -log(0.2)\n",
        "# P(Grippy)=0.8  -> E1 = -log(0.8)\n",
        "scalar0 = -np.log(0.2)\n",
        "scalar1 = -np.log(0.8)\n",
        "hnlf_b = HybridNonlinearFactor(dk0, [(odom0, scalar0), (odom1, scalar1)])\n",
        "print(\"\\nHybridNonlinearFactor (factors + scalars):\")\n",
        "hnlf_b.print()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Accessing Components and Error Calculation\n",
        "\n",
        "Similar to [`HybridGaussianFactor`](HybridGaussianFactor.ipynb), you can access the underlying decision tree and calculate errors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Error for D0=0 (Slip): 1.6694379124341003\n",
            "Error using HybridValues (Slip): 1.6694379124341003\n",
            "\n",
            "Error for D0=1 (Grippy): 3.2231435513142106\n",
            "\n",
            "Error Tree Choice(d0) \n",
            "Error Tree 0 Leaf 1.6694379\n",
            "Error Tree 1 Leaf 3.2231436\n"
          ]
        }
      ],
      "source": [
        "# --- Error Calculation ---\n",
        "# Requires continuous Values and discrete DiscreteValues (or HybridValues)\n",
        "\n",
        "# Continuous values\n",
        "cont_vals = Values()\n",
        "cont_vals.insert(X(0), Pose2(0, 0, 0))\n",
        "cont_vals.insert(X(1), Pose2(1.1, 0.1, np.radians(2))) # Slightly off\n",
        "\n",
        "# Discrete assignment\n",
        "assignment0 = DiscreteValues()\n",
        "assignment0[D(0)] = 0  # Slip mode\n",
        "\n",
        "# Method 1: error(Values, DiscreteValues)\n",
        "err0 = hnlf_b.error(cont_vals, assignment0)\n",
        "# Expected: scalar0 + 0.5 * || odom0.evaluateError(x0, x1) ||^2\n",
        "print(f\"\\nError for D0=0 (Slip): {err0}\")\n",
        "\n",
        "# Method 2: error(HybridValues)\n",
        "hybrid_vals0 = HybridValues(cv=gtsam.VectorValues(), dv=assignment0, v=cont_vals)\n",
        "err0_hv = hnlf_b.error(hybrid_vals0)\n",
        "print(f\"Error using HybridValues (Slip): {err0_hv}\")\n",
        "\n",
        "# Check Grippy mode\n",
        "assignment1 = DiscreteValues()\n",
        "assignment1[D(0)] = 1  # Grippy mode\n",
        "err1 = hnlf_b.error(cont_vals, assignment1)\n",
        "print(f\"\\nError for D0=1 (Grippy): {err1}\\n\") # Should be much lower due to smaller noise\n",
        "\n",
        "# --- errorTree(Values) ---\n",
        "# Calculates nonlinear error for *all* discrete modes given continuous values\n",
        "adt = hnlf_b.errorTree(cont_vals)\n",
        "adt.print(\"Error Tree\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Linearization\n",
        "\n",
        "A key function is `linearize`, which converts the `HybridNonlinearFactor` into a `HybridGaussianFactor` at a given linearization point (continuous `Values`)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Linearized HybridGaussianFactor:\n",
            "HybridGaussianFactor\n",
            "\n",
            "Hybrid [x0 x1; d0]{\n",
            " Choice(d0) \n",
            " 0 Leaf :\n",
            "  A[x0] = [\n",
            "\t-2, -0, -0;\n",
            "\t0, -2, -2;\n",
            "\t-0, -0, -5.72957795\n",
            "]\n",
            "  A[x1] = [\n",
            "\t2, 0, 0;\n",
            "\t0, 2, 0;\n",
            "\t0, 0, 5.72957795\n",
            "]\n",
            "  b = [ -0 -0 -0 ]\n",
            "  No noise model\n",
            "scalar: 1.23431728\n",
            "\n",
            " 1 Leaf :\n",
            "  A[x0] = [\n",
            "\t-10, -0, -0;\n",
            "\t0, -10, -10;\n",
            "\t-0, -0, -57.2957795\n",
            "]\n",
            "  A[x1] = [\n",
            "\t10, 0, 0;\n",
            "\t0, 10, 0;\n",
            "\t0, 0, 57.2957795\n",
            "]\n",
            "  b = [ -0 -0 -0 ]\n",
            "  No noise model\n",
            "scalar: -5.673438\n",
            "\n",
            "}\n"
          ]
        }
      ],
      "source": [
        "# Using hnlf_b\n",
        "linearization_point = Values()\n",
        "linearization_point.insert(X(0), Pose2(0, 0, 0))\n",
        "linearization_point.insert(X(1), Pose2(1.0, 0, 0)) # Linearize at the expected mean\n",
        "\n",
        "# Linearize the whole factor (all modes)\n",
        "hybrid_gaussian_factor = hnlf_b.linearize(linearization_point)\n",
        "print(\"\\nLinearized HybridGaussianFactor:\")\n",
        "hybrid_gaussian_factor.print()\n",
        "# Note: The Gaussian components will be JacobianFactors evaluated at the linearization point.\n",
        "# The scalar energies E_m are carried over.\n",
        "\n",
        "# Linearize only for a specific mode (useful internally, maybe not direct API)\n",
        "# assignment = gtsam.DiscreteValues([(D(0), 0)]) # Slippery\n",
        "# gaussian_factor_mode0 = hnlf_b.linearize(linearization_point, assignment) # Hypothetical API\n",
        "# print(\"\\nLinearized GaussianFactor for Mode 0:\")\n",
        "# gaussian_factor_mode0.print()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Restriction (`restrict`)\n",
        "\n",
        "Fixes the discrete variables, returning the corresponding `NoiseModelFactor` component."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Restricted Factor for D0=1 (Grippy):BetweenFactor(x0,x1)\n",
            "  measured:  (1, 0, 0)\n",
            "  noise model: diagonal sigmas [0.1; 0.1; 0.0174532925];\n"
          ]
        }
      ],
      "source": [
        "# Using hnlf_b\n",
        "# Create an empty DiscreteValues and assign D(0) = 1 (Grippy mode)\n",
        "assignment = DiscreteValues()\n",
        "assignment[D(0)] = 1\n",
        "restricted_factor_ptr = hnlf_b.restrict(assignment)\n",
        "\n",
        "restricted_factor_ptr.print(\"\\nRestricted Factor for D0=1 (Grippy):\")"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "py312",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.6"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}
